import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from ocr import OcrTableAccurate
from Chat import ChatCompletions
import os
import numpy as np
from datetime import datetime

class ExcelAnalyzer:
    def __init__(self, secret_id, secret_key):
        self.secret_id = secret_id
        self.secret_key = secret_key
        self.ocr = OcrTableAccurate(secret_id, secret_key)
        self.chat = ChatCompletions(secret_id, secret_key)
        self.current_df = None
        
        # 设置matplotlib中文支持
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 设置seaborn样式
        sns.set_theme(font='SimHei', font_scale=1.2)
        
        # 定义可用的图表类型
        self.available_charts = {
            "折线图": self.plot_line,
            "柱状图": self.plot_bar,
            "散点图": self.plot_scatter,
            "饼图": self.plot_pie,
            "箱线图": self.plot_box,
            "直方图": self.plot_hist,
            "热力图": self.plot_heatmap,
            "多子图": self.plot_subplots,
            "堆叠柱状图": self.plot_stacked_bar,
            "面积图": self.plot_area,
            "气泡图": self.plot_bubble,
            "3D散点图": self.plot_3d_scatter
        }

    def load_data(self, file_path):
        """加载数据文件"""
        try:
            file_ext = os.path.splitext(file_path)[1].lower()
            
            # 处理图片文件
            if file_ext in ['.png', '.jpg', '.jpeg', '.bmp', '.gif']:
                excel_path = os.path.splitext(file_path)[0] + "_ocr.xlsx"
                success = self.image_to_excel(file_path, excel_path)
                if not success:
                    return None, "图片OCR转换失败"
                file_path = excel_path
            
            self.current_df = pd.read_excel(file_path)
            return self.current_df, "数据加载成功"
            
        except Exception as e:
            return None, f"数据加载失败: {str(e)}"

    def get_data_preview(self):
        """获取数据预览信息"""
        if self.current_df is None:
            return "请先加载数据"
            
        preview = "##  数据概览\n\n"
        # ... (保持原有的预览生成代码)
        return preview

    def analyze_data(self):
        """生成数据分析报告"""
        try:
            if self.current_df is None or self.current_df.empty:
                return "请先上传数据文件"
            
            report = "##  1 数据分析报告\n\n"
            
            # 1. 基本信息
            report += "### 1️ 基本信息\n"
            report += f"- 总行数：{len(self.current_df)}\n"
            report += f"- 总列数：{len(self.current_df.columns)}\n"
            report += f"- 数据维度：{self.current_df.shape}\n\n"
            
            # 2. 数据类型分析
            report += "### 2️ 数据类型分析\n"
            report += "```\n"
            report += self.current_df.dtypes.to_string()
            report += "\n```\n\n"
            
            # 3. 基本统计分析
            report += "### 3️ 基本统计分析\n"
            report += self.current_df.describe().to_markdown()
            report += "\n\n"
            
            # 4. 缺失值分析
            missing = self.current_df.isnull().sum()
            if missing.any():
                report += "### 4 缺失值分析\n"
                report += "```\n"
                report += missing.to_string()
                report += "\n```\n\n"
            
            # 5. 数值列分析
            numeric_cols = self.current_df.select_dtypes(include=['number']).columns
            if len(numeric_cols) > 0:
                report += "### 5 数值列分析\n"
                for col in numeric_cols:
                    report += f"#### {col}\n"
                    stats = self.current_df[col].describe()
                    report += f"- 均值：{stats['mean']:.2f}\n"
                    report += f"- 中位数：{stats['50%']:.2f}\n"
                    report += f"- 标准差：{stats['std']:.2f}\n"
                    report += f"- 偏度：{self.current_df[col].skew():.2f}\n"
                    report += f"- 峰度：{self.current_df[col].kurtosis():.2f}\n"
                    report += f"- 最小值：{stats['min']:.2f}\n"
                    report += f"- 最大值：{stats['max']:.2f}\n\n"
            
            # 6. 分类列分析
            categorical_cols = self.current_df.select_dtypes(include=['object']).columns
            if len(categorical_cols) > 0:
                report += "### 6️ 分类列分析\n"
                for col in categorical_cols:
                    report += f"#### {col}\n"
                    value_counts = self.current_df[col].value_counts()
                    report += "```\n"
                    report += value_counts.to_string()
                    report += "\n```\n\n"
            
            # 7. 相关性分析
            if len(numeric_cols) > 1:
                report += "### 7️ 相关性分析\n"
                corr = self.current_df[numeric_cols].corr()
                report += corr.to_markdown()
                report += "\n\n"
            
            return report
        except Exception as e:
            return f"分析失败: {str(e)}"

    def chat_with_data(self, message):
        """与数据对话"""
        if self.current_df is None:
            return "请先加载数据"
            
        # 构建数据上下文
        data_context = self._build_data_context(message)
        # 调用大模型分析
        response = self.chat.chat_completions(data_context)
        return self._parse_response(response)

    def generate_chart(self, chart_type, x_col, y_col, z_col=None, title=""):
        """生成图表"""
        try:
            if self.current_df is None or self.current_df.empty:
                return "请先上传数据", None
            
            if not x_col or not y_col:
                return "请选择X轴和Y轴数据", None
            
            if chart_type in ["气泡图", "3D散点图"] and not z_col:
                return "请选择Z轴/大小数据", None
            
            # 检查列是否存在
            if x_col not in self.current_df.columns or y_col not in self.current_df.columns:
                return "选择的列不存在", None
            if z_col and z_col not in self.current_df.columns:
                return "Z轴/大小数据列不存在", None
            
            # 生成图表
            fig_path = self.available_charts[chart_type](
                x_col, y_col, z_col, title,
                figsize=(10, 6),
                fontsize=12,
                rotation=45,
                grid=True,
                grid_alpha=0.3
            )
            
            return "图表生成成功", fig_path
            
        except Exception as e:
            return f"图表生成失败: {str(e)}", None

    def plot_line(self, x_col, y_col, z_col=None, title="", **kwargs):
        """绘制折线图"""
        plt.figure(figsize=kwargs.get('figsize', (10, 6)))
        plt.plot(self.current_df[x_col], self.current_df[y_col], 
                marker='o', linestyle='-', linewidth=2)
        
        plt.xlabel(x_col, fontproperties='SimHei', fontsize=12)
        plt.ylabel(y_col, fontproperties='SimHei', fontsize=12)
        if title:
            plt.title(title, fontproperties='SimHei', fontsize=14)
        plt.grid(True, alpha=0.3)
        plt.xticks(fontproperties='SimHei', rotation=45)
        plt.yticks(fontproperties='SimHei')
        
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def plot_bar(self, x_col, y_col, z_col=None, title="", **kwargs):
        """绘制柱状图"""
        plt.figure(figsize=kwargs.get('figsize', (10, 6)))
        if y_col:
            data = self.current_df.groupby(x_col)[y_col].mean()
            data.plot(kind='bar', color=kwargs.get('color', 'skyblue'))
            plt.ylabel(y_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        else:
            self.current_df[x_col].value_counts().plot(kind='bar', color=kwargs.get('color', 'skyblue'))
        plt.xlabel(x_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        if title:
            plt.title(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12)+2)
        if kwargs.get('grid', True):
            plt.grid(True, alpha=kwargs.get('grid_alpha', 0.3))
        plt.xticks(fontproperties='SimHei', rotation=kwargs.get('rotation', 45))
        plt.yticks(fontproperties='SimHei')
        
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def plot_scatter(self, x_col, y_col, z_col=None, title="", **kwargs):
        """绘制散点图"""
        plt.figure(figsize=kwargs.get('figsize', (10, 6)))
        plt.scatter(
            self.current_df[x_col], 
            self.current_df[y_col],
            alpha=0.6,
            color=kwargs.get('color', 'skyblue')
        )
        
        plt.xlabel(x_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        plt.ylabel(y_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        if title:
            plt.title(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12)+2)
        if kwargs.get('grid', True):
            plt.grid(True, alpha=kwargs.get('grid_alpha', 0.3))
        plt.xticks(fontproperties='SimHei', rotation=kwargs.get('rotation', 45))
        plt.yticks(fontproperties='SimHei')
        
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def plot_pie(self, x_col, y_col=None, z_col=None, title="", **kwargs):
        """绘制饼图"""
        plt.figure(figsize=kwargs.get('figsize', (10, 8)))
        data = self.current_df[x_col].value_counts()
        colors = plt.cm.Pastel1(np.linspace(0, 1, len(data)))
        plt.pie(
            data, 
            labels=data.index, 
            autopct='%1.1f%%', 
            colors=colors,
            textprops={
                'fontproperties': 'SimHei', 
                'fontsize': kwargs.get('fontsize', 12)
            }
        )
        if title:
            plt.title(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12)+2)
        
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def plot_box(self, x_col, y_col=None, z_col=None, title="", **kwargs):
        """绘制箱线图"""
        plt.figure(figsize=kwargs.get('figsize', (10, 6)))
        if y_col:
            sns.boxplot(
                x=self.current_df[x_col], 
                y=self.current_df[y_col],
                color=kwargs.get('color', 'skyblue')
            )
            plt.xlabel(x_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
            plt.ylabel(y_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        else:
            sns.boxplot(
                y=self.current_df[x_col],
                color=kwargs.get('color', 'skyblue')
            )
            plt.ylabel(x_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        
        if title:
            plt.title(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12)+2)
        if kwargs.get('grid', True):
            plt.grid(True, alpha=kwargs.get('grid_alpha', 0.3))
        plt.xticks(fontproperties='SimHei', rotation=kwargs.get('rotation', 45))
        plt.yticks(fontproperties='SimHei')
        
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def plot_hist(self, x_col, y_col=None, z_col=None, title="", **kwargs):
        """绘制直方图"""
        plt.figure(figsize=kwargs.get('figsize', (10, 6)))
        plt.hist(
            self.current_df[x_col], 
            bins=30, 
            alpha=0.7, 
            color=kwargs.get('color', 'skyblue'),
            edgecolor='black'
        )
        plt.xlabel(x_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        plt.ylabel('频率', fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        if title:
            plt.title(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12)+2)
        if kwargs.get('grid', True):
            plt.grid(True, alpha=kwargs.get('grid_alpha', 0.3))
        plt.xticks(fontproperties='SimHei', rotation=kwargs.get('rotation', 45))
        plt.yticks(fontproperties='SimHei')
        
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def plot_heatmap(self, x_col, y_col=None, z_col=None, title="", **kwargs):
        """绘制热力图"""
        plt.figure(figsize=kwargs.get('figsize', (10, 8)))
        
        # 计算相关性矩阵
        numeric_cols = self.current_df.select_dtypes(include=['number']).columns
        corr = self.current_df[numeric_cols].corr()
        
        # 绘制热力图
        sns.heatmap(
            corr,
            annot=True,  # 显示数值
            fmt='.2f',   # 数值格式
            cmap='coolwarm',  # 色彩方案
            square=True,  # 正方形单元格
            annot_kws={'size': kwargs.get('fontsize', 10)}
        )
        
        if title:
            plt.title(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12)+2)
        plt.xticks(fontproperties='SimHei', rotation=45)
        plt.yticks(fontproperties='SimHei', rotation=45)
        
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def plot_subplots(self, x_col, y_col, z_col=None, title="", **kwargs):
        """绘制多子图"""
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        
        # 折线图
        axes[0,0].plot(self.current_df[x_col], self.current_df[y_col])
        axes[0,0].set_title('折线图', fontproperties='SimHei')
        axes[0,0].set_xlabel(x_col, fontproperties='SimHei')
        axes[0,0].set_ylabel(y_col, fontproperties='SimHei')
        
        # 散点图
        axes[0,1].scatter(self.current_df[x_col], self.current_df[y_col])
        axes[0,1].set_title('散点图', fontproperties='SimHei')
        axes[0,1].set_xlabel(x_col, fontproperties='SimHei')
        axes[0,1].set_ylabel(y_col, fontproperties='SimHei')
        
        # 柱状图
        axes[1,0].bar(self.current_df[x_col], self.current_df[y_col])
        axes[1,0].set_title('柱状图', fontproperties='SimHei')
        axes[1,0].set_xlabel(x_col, fontproperties='SimHei')
        axes[1,0].set_ylabel(y_col, fontproperties='SimHei')
        
        # 箱线图
        sns.boxplot(data=self.current_df, y=y_col, ax=axes[1,1])
        axes[1,1].set_title('箱线图', fontproperties='SimHei')
        
        if title:
            fig.suptitle(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 14))
        
        plt.tight_layout()
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def plot_stacked_bar(self, x_col, y_col, z_col=None, title="", **kwargs):
        """绘制堆叠柱状图"""
        plt.figure(figsize=kwargs.get('figsize', (10, 6)))
        
        # 对数据进行分组和堆叠
        if y_col:
            data = self.current_df.pivot_table(
                index=x_col, 
                columns=y_col, 
                aggfunc='size',
                fill_value=0
            )
            data.plot(kind='bar', stacked=True)
            plt.ylabel('数量', fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        else:
            self.current_df[x_col].value_counts().plot(kind='bar')
        
        plt.xlabel(x_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        if title:
            plt.title(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12)+2)
        if kwargs.get('grid', True):
            plt.grid(True, alpha=kwargs.get('grid_alpha', 0.3))
        plt.xticks(rotation=kwargs.get('rotation', 45), fontproperties='SimHei')
        plt.yticks(fontproperties='SimHei')
        plt.legend(prop={'family': 'SimHei'})
        
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def plot_area(self, x_col, y_col, z_col=None, title="", **kwargs):
        """绘制面积图"""
        plt.figure(figsize=kwargs.get('figsize', (10, 6)))
        
        plt.fill_between(
            self.current_df[x_col], 
            self.current_df[y_col], 
            alpha=0.5,
            color=kwargs.get('color', 'skyblue')
        )
        plt.plot(self.current_df[x_col], self.current_df[y_col], 
                color=kwargs.get('color', 'skyblue'))
        
        plt.xlabel(x_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        plt.ylabel(y_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        if title:
            plt.title(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12)+2)
        if kwargs.get('grid', True):
            plt.grid(True, alpha=kwargs.get('grid_alpha', 0.3))
        plt.xticks(rotation=kwargs.get('rotation', 45), fontproperties='SimHei')
        plt.yticks(fontproperties='SimHei')
        
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def plot_bubble(self, x_col, y_col, z_col, title="", **kwargs):
        """绘制气泡图"""
        plt.figure(figsize=kwargs.get('figsize', (10, 6)))
        
        # 使用z_col作为气泡大小
        sizes = self.current_df[z_col] * 100  # 调整气泡大小
        
        plt.scatter(
            self.current_df[x_col], 
            self.current_df[y_col],
            s=sizes,
            alpha=0.6,
            color=kwargs.get('color', 'skyblue')
        )
        
        plt.xlabel(x_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        plt.ylabel(y_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        if title:
            plt.title(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12)+2)
        if kwargs.get('grid', True):
            plt.grid(True, alpha=kwargs.get('grid_alpha', 0.3))
        plt.xticks(rotation=kwargs.get('rotation', 45), fontproperties='SimHei')
        plt.yticks(fontproperties='SimHei')
        
        # 添加图例说明Z轴数据
        plt.colorbar(label=z_col)
        
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path

    def _build_data_context(self, message):
        """构建数据上下文"""
        basic_stats = {
            'total_rows': self.current_df.shape[0],
            'total_cols': self.current_df.shape[1],
            'null_count': self.current_df.isnull().sum().sum(),
            'numeric_cols': self.current_df.select_dtypes(include=['int64', 'float64']).columns.tolist(),
            'categorical_cols': self.current_df.select_dtypes(include=['object']).columns.tolist()
        }
        
        return f"""
        数据分析上下文：
        1. 数据概况：
           - 数据规模：{basic_stats['total_rows']}行 × {basic_stats['total_cols']}列
           - 缺失值数量：{basic_stats['null_count']}
           - 数值型列：{', '.join(basic_stats['numeric_cols'])}
           - 类别型列：{', '.join(basic_stats['categorical_cols'])}
        
        2. 数据内容：
           {self.current_df.to_string()}
        
        用户问题：{message}
        """

    def _set_chart_style(self, style_params):
        """设置图表样式"""
        if style_params.get('color_theme') == "深色":
            plt.style.use('dark_background')
        elif style_params.get('color_theme') == "浅色":
            plt.style.use('seaborn-pastel')
        else:
            plt.style.use('default')

    def image_to_excel(self, image_path, output_excel_path):
        """将图片转换为Excel文件"""
        try:
            # OCR识别表格
            result = self.ocr.recognize_table_accurate(image_path)
            
            # 解析OCR结果
            if 'TableDetections' in result:
                table_data = result['TableDetections'][0]['Cells']  # 获取单元格数据
                
                # 找出表格的最大行列数
                max_row = max(cell['RowBr'] for cell in table_data)
                max_col = max(cell['ColBr'] for cell in table_data)
                
                # 创建空的二维数组
                rows = [['' for _ in range(max_col)] for _ in range(max_row)]
                
                # 填充数据
                for cell in table_data:
                    row_start = cell['RowTl']
                    col_start = cell['ColTl']
                    text = cell['Text']
                    # 处理合并单元格
                    row_end = cell['RowBr']
                    col_end = cell['ColBr']
                    
                    # 填充所有涉及的单元格
                    for r in range(row_start, row_end):
                        for c in range(col_start, col_end):
                            rows[r][c] = text
                
                # 创建DataFrame
                df = pd.DataFrame(rows)
                
                # 使用第一行作为列名
                df.columns = df.iloc[0]
                df = df.iloc[1:]  # 删除第一行（现在是列名了）
                
                # 清理数据
                # 删除全为空的行和列
                df = df.dropna(how='all')
                df = df.dropna(axis=1, how='all')
                
                # 重置索引
                df = df.reset_index(drop=True)
                
                # 保存为Excel
                df.to_excel(output_excel_path, index=False)
                return True
                
        except Exception as e:
            print(f"转换失败: {str(e)}")
            return False

    def plot_3d_scatter(self, x_col, y_col, z_col, title="", **kwargs):
        """绘制3D散点图"""
        from mpl_toolkits.mplot3d import Axes3D  # 导入3D工具包
        
        fig = plt.figure(figsize=kwargs.get('figsize', (10, 8)))
        ax = fig.add_subplot(111, projection='3d')  # 创建3D子图
        
        # 绘制散点
        scatter = ax.scatter(
            self.current_df[x_col],
            self.current_df[y_col],
            self.current_df[z_col],
            c=self.current_df[z_col],  # 使用Z值着色
            cmap='viridis',  # 色彩方案
            alpha=0.6
        )
        
        # 设置轴标签
        ax.set_xlabel(x_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        ax.set_ylabel(y_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        ax.set_zlabel(z_col, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12))
        
        # 设置标题
        if title:
            ax.set_title(title, fontproperties='SimHei', fontsize=kwargs.get('fontsize', 12)+2)
        
        # 添加颜色条
        plt.colorbar(scatter, label=z_col)
        
        # 设置网格
        if kwargs.get('grid', True):
            ax.grid(True, alpha=kwargs.get('grid_alpha', 0.3))
        
        # 调整视角
        ax.view_init(elev=20, azim=45)  # 设置观察角度
        
        # 保存图片
        fig_path = "temp_chart.png"
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        plt.close()
        return fig_path 